import logging

import torch
from torch import Tensor, nn
import torch.nn.functional as F

import torch


def dist_loss(input: Tensor):
    # input的形状为(batch_size, points_num, 3)
    # 只使用前两个维度的坐标信息
    coordinates = input[:, :, :2]  # (batch_size, points_num, 2)

    batch_size, points_num, _ = coordinates.size()

    # 计算所有点对的欧式距离
    # 展开coordinates以计算点对之间的差值
    expanded_coords = coordinates.unsqueeze(2)  # (batch_size, points_num, 1, 2)
    # 计算点对的差值
    diffs = expanded_coords - expanded_coords.transpose(1, 2)  # (batch_size, points_num, points_num, 2)
    # 计算欧式距离的平方
    dists_squared = (diffs ** 2).sum(dim=3)  # (batch_size, points_num, points_num)

    # 添加一个小的epsilon防止除以0
    epsilon = 1e-8
    # 计算距离的倒数，忽略对角线上的元素（点到其自身的距离）
    dists_squared += torch.eye(points_num).to(dists_squared.device) * epsilon  # 防止自距离为0
    inv_dists = 1.0 / dists_squared

    # 创建一个掩码来忽略对角线上的元素
    mask = torch.eye(points_num).unsqueeze(0).to(inv_dists.device) == 0
    max_inv_dists = inv_dists.masked_fill(~mask, 0).view(mask.shape[0], -1)
    max_inv_dists, indices = max_inv_dists.topk(5, dim=1)

    # 我们期望点之间的距离尽可能大，即这个损失越小越好
    loss = max_inv_dists.mean()  # 对所有batch的损失取平均

    logging.info('dist loss: {}'.format(loss))
    return loss


def balance_loss(input: Tensor):
    # 获取batch_size和points_num
    batch_size, points_num, _ = input.size()

    # 获取标签置信度，形状为(batch_size, points_num)
    logits = input[:, :, 2]

    labels = torch.sigmoid(logits)

    # 将标签置信度从[0,1]转换到[-0.5, 0.5]区间
    shifted_labels = labels - 0.5

    # 求和所有负数值和所有正数值
    negative_sum = torch.where(shifted_labels < 0, shifted_labels, torch.zeros_like(shifted_labels)).sum(dim=1)
    positive_sum = torch.where(shifted_labels > 0, shifted_labels, torch.zeros_like(shifted_labels)).sum(dim=1)

    # 计算损失，这里可以使用mean而非sum来避免数值过大
    loss = torch.mean(torch.abs(negative_sum + positive_sum))
    logging.info('balance loss: {}'.format(loss))
    return loss


def label_accuracy_loss(input: torch.Tensor, mask: torch.Tensor):
    batch_size, points_num, _ = input.size()
    _, height, width = mask.size()

    # 获取点的坐标和置信度
    coords = input[:, :, :2]
    coords = torch.sigmoid(coords)

    logits = input[:, :, 2]

    # 将坐标转换为像素坐标
    x_coords = (coords[:, :, 0] * width).long()
    y_coords = (coords[:, :, 1] * height).long()

    # 防止坐标超出图像边界
    x_coords = torch.clamp(x_coords, 0, width - 1)
    y_coords = torch.clamp(y_coords, 0, height - 1)

    # 使用像素坐标从mask中提取值
    mask_values = mask[torch.arange(batch_size)[:, None], y_coords, x_coords]

    # 计算损失
    loss_fn = nn.BCEWithLogitsLoss()
    loss = loss_fn(logits, mask_values.float())

    logging.info('label_accuracy_loss: {}'.format(loss))
    return loss


def get_loss(masks_pred: torch.Tensor, true_masks: torch.Tensor):
    true_masks = 1 - true_masks
    return 0.001 * dist_loss(masks_pred), 2 * balance_loss(masks_pred), 3 * label_accuracy_loss(masks_pred, true_masks)
